{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.102704Z",
     "start_time": "2025-05-25T09:36:41.093857Z"
    },
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[813, 658, 771, 410, 517],\n",
       "        [180, 995, 954, 174, 953]])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "batch_size = 2\n",
    "seq_len = 5\n",
    "embedding_size = 100\n",
    "vocab_size = 1000\n",
    "\n",
    "# 代表输入sequence\n",
    "# 随机生成0-999的数字，形状为(batch_size, seq_len)，表示batch_size条sequence，每条sequence的长度为seq_len\n",
    "inputs = torch.randint(0, vocab_size, (batch_size, seq_len))\n",
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "d75a5f086f4eea9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.168826Z",
     "start_time": "2025-05-25T09:36:41.165040Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# input_embedding层\n",
    "# i表示输入\n",
    "input_embedding = nn.Embedding(vocab_size, embedding_size)\n",
    "input_embeddings = input_embedding(inputs)\n",
    "input_embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "b78fcbd2c1eeab9e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.226052Z",
     "start_time": "2025-05-25T09:36:41.221810Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 100])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 位置编码\n",
    "# 位置编码只跟位置有关，和具体位置上输入的数据无关\n",
    "pe = torch.zeros(seq_len, embedding_size)\n",
    "pos = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)\n",
    "pe[:, 0::2] = torch.sin(pos / (10000 ** (2 * torch.arange(0, embedding_size, 2) / embedding_size)))\n",
    "pe[:, 1::2] = torch.cos(pos / (10000 ** (2 * torch.arange(1, embedding_size, 2) / embedding_size)))\n",
    "pe.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "220f121c2cc2e2ef",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.264271Z",
     "start_time": "2025-05-25T09:36:41.261263Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 位置编码+input_embeddings\n",
    "inputs_encoder = input_embeddings + pe\n",
    "inputs_encoder.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "b87ee8e199a33508",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.290990Z",
     "start_time": "2025-05-25T09:36:41.284618Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Multi-Head Attention多头注意力\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, embed_dim, head_dim, num_heads):\n",
    "        super().__init__()\n",
    "\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = head_dim\n",
    "\n",
    "        # 合并所有头的线性变换矩阵\n",
    "        self.q_proj = nn.Linear(embed_dim, head_dim * num_heads)  # 维度(3, 4)\n",
    "        self.k_proj = nn.Linear(embed_dim, head_dim * num_heads)\n",
    "        self.v_proj = nn.Linear(embed_dim, head_dim * num_heads)\n",
    "        self.out_proj = nn.Linear(head_dim * num_heads, embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "\n",
    "        # 计算所有头的qkv\n",
    "        q = self.q_proj(x)  # x是(4, 3), 输出是(4, 4)\n",
    "        k = self.k_proj(x)  # x是(4, 3), 输出是(4, 4)\n",
    "        v = self.v_proj(x)  # x是(4, 3), 输出是(4, 4)\n",
    "\n",
    "        # 分割为多个头\n",
    "        # view()的作用是将张量reshape为指定形状，也就是(4, 2, 2)\n",
    "        # transpose(0, 1)的作用是将维度1和维度2进行转置，也就是(2, 4, 2)把num_heads放到前面去，这样就方便去每个head的qkv\n",
    "        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # 计算注意力得分 [num_heads, seq_len, seq_len]，最终形状为(2, 4, 4)\n",
    "        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)  #添加缩放因子，防止乘积过大\n",
    "\n",
    "        # 计算注意力权重 [num_heads, seq_len, seq_len]，最终形状为(2, 4, 4)\n",
    "        attn_weights = torch.softmax(scores, dim=-1)\n",
    "\n",
    "        # 应用注意力权重 [num_heads, seq_len, head_dim]，最终形状为(2, 4, 2)\n",
    "        output = torch.matmul(attn_weights, v)\n",
    "\n",
    "        # 合并所有头\n",
    "        # output.transpose(0, 1) [batch_size, num_heads, seq_len, head_dim] -> [seq_len, num_heads, head_dim]\n",
    "        output = output.transpose(0, 1).contiguous()\n",
    "\n",
    "        # [seq_len, num_heads, head_dim] -> [seq_len, num_heads * head_dim] 也就是(4, 2*2)\n",
    "        output = output.view(batch_size, seq_len, -1)\n",
    "\n",
    "        # 最终输出投影\n",
    "        return self.out_proj(output)\n",
    "\n",
    "\n",
    "mha = MultiHeadAttention(embed_dim=embedding_size, head_dim=2, num_heads=2)\n",
    "\n",
    "outputs_attention = mha(inputs_encoder)\n",
    "\n",
    "outputs_attention.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "7cadc119cd37aaa3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.317012Z",
     "start_time": "2025-05-25T09:36:41.314097Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 残差连接\n",
    "outputs_attention = outputs_attention + inputs_encoder\n",
    "outputs_attention.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "fa842128f70b648e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.349957Z",
     "start_time": "2025-05-25T09:36:41.346937Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 层归一化\n",
    "layer_norm = nn.LayerNorm(embedding_size)\n",
    "outputs_attention = layer_norm(outputs_attention)\n",
    "outputs_attention.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "a004b60ec103d66a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.386732Z",
     "start_time": "2025-05-25T09:36:41.382195Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Feed Forward，两个线程层，最后输出维度不变\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, embed_dim, fc_dim):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(embed_dim, fc_dim)\n",
    "        self.fc2 = nn.Linear(fc_dim, embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc2(torch.relu(self.fc1(x)))\n",
    "\n",
    "feed_forward = FeedForward(embed_dim=embedding_size, fc_dim=embedding_size * 4)\n",
    "outputs_encoder = feed_forward(outputs_attention)\n",
    "outputs_encoder.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "1d01b7ff51008594",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-25T09:36:41.415380Z",
     "start_time": "2025-05-25T09:36:41.412114Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 再来一次残差连接和层归一化，得到最终的Encoder输出\n",
    "outputs_encoder = layer_norm(outputs_encoder + outputs_attention)\n",
    "outputs_encoder.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "16fba6d2199cee2b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-26T00:57:58.171243Z",
     "start_time": "2025-05-26T00:57:58.151007Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 定义Encoder\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, embed_dim, fc_dim, num_heads, num_layers):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.layers = nn.ModuleList([\n",
    "            EncoderLayer(embed_dim, fc_dim, num_heads)\n",
    "            for _ in range(num_layers)\n",
    "        ])\n",
    "\n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return x\n",
    "\n",
    "# 定义EncoderLayer\n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, embed_dim, fc_dim, num_heads):\n",
    "        super().__init__()\n",
    "        self.mha = MultiHeadAttention(embed_dim, embed_dim // num_heads, num_heads)\n",
    "        self.feed_forward = FeedForward(embed_dim, fc_dim)\n",
    "        self.layer_norm1 = nn.LayerNorm(embed_dim)\n",
    "        self.layer_norm2 = nn.LayerNorm(embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 1.多头注意力\n",
    "        # 2.残差连接和层归一化\n",
    "        x = self.layer_norm1(x + self.mha(x))\n",
    "\n",
    "        # 3.Feed Forward\n",
    "        # 4.残差连接和层归一化\n",
    "        return self.layer_norm2(x + self.feed_forward(x))\n",
    "\n",
    "encoder = Encoder(embed_dim=embedding_size, fc_dim=embedding_size * 4, num_heads=2, num_layers=2)\n",
    "\n",
    "encoder_outputs = encoder(inputs_encoder)\n",
    "encoder_outputs.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
